package com.lebang.medicalInquiry.component;


import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.time.Instant;
import java.util.Arrays;
import java.util.List;

/**
 * 频率限制注解
 */
@Component
@Slf4j
public class RateLimitComponent {

    /**
     * 默认使用的 redis
     */
    @Resource(name = "defaultRedisTemplate")
    private RedisTemplate<String, Object> redisTemplate;

    private DefaultRedisScript script;

    @PostConstruct
    public void init() {
        this.script = new DefaultRedisScript<List>();
        script.setScriptSource(new ResourceScriptSource(new ClassPathResource("scripts/request_rate_limiter.lua")));
        script.setResultType(List.class);
    }

    public boolean isAllow(String key, int replenishRate, int burstCapacity, int requestedTokens) {
        try {
            List<String> keys = getKeys(key);
            List execute = (List) this.redisTemplate.execute(this.script, keys, replenishRate,
                    burstCapacity, Instant.now().getEpochSecond(),
                    requestedTokens);
            return Integer.parseInt(execute.get(0) + "") == 1;
        } catch (Exception e) {
            log.error("限流异常");
            e.printStackTrace();
            return false;
        }
    }


    List<String> getKeys(String id) {
        String prefix = "rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }
}
